package scanner;

import org.apache.commons.io.FileUtils;
import org.eclipse.jdt.core.JavaCore;
import org.eclipse.jdt.core.dom.AST;
import org.eclipse.jdt.core.dom.ASTParser;
import org.eclipse.jdt.core.dom.ASTVisitor;
import org.eclipse.jdt.core.dom.CompilationUnit;
import org.eclipse.jdt.core.dom.IDocElement;
import org.eclipse.jdt.core.dom.Javadoc;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.PackageDeclaration;
import org.eclipse.jdt.core.dom.SingleVariableDeclaration;
import org.eclipse.jdt.core.dom.TagElement;
import org.eclipse.jdt.core.dom.ThrowStatement;
import org.eclipse.jdt.core.dom.Type;
import org.eclipse.jdt.core.dom.TypeDeclaration;
import org.eclipse.jdt.core.dom.TypeParameter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.Stack;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * Created by Shunjie Ding on 2018/5/26.
 */
public class UncheckedExceptionSourceScanner {
    private static final String DEFAULT_ENCODING = "UTF-8";
    private static final Charset DEFAULT_CHARSET = Charset.forName(DEFAULT_ENCODING);
    private Logger logger = LoggerFactory.getLogger(UncheckedExceptionSourceScanner.class);
    private String[] classpath;

    public UncheckedExceptionSourceScanner(String[] classpath) {
        this.classpath = classpath;
    }

    private static String readFileToString(File file) throws IOException {
        return FileUtils.readFileToString(file, DEFAULT_CHARSET);
    }

    CompilationUnit compileSource(String[] sources, String sourceFile, String sourceRoot)
        throws IOException {
        return compileSource(sources, new File(sourceFile), new File(sourceFile));
    }

    CompilationUnit compileSource(String[] sources, File sourceFile, File sourceRoot)
        throws IOException {
        ASTParser parser = ASTParser.newParser(AST.JLS10);
        // setting up ast parser
        parser.setEnvironment(classpath, sources, new String[] {DEFAULT_ENCODING}, true);

        parser.setResolveBindings(true);
        parser.setBindingsRecovery(true);
        parser.setStatementsRecovery(true);
        parser.setCompilerOptions(JavaCore.getOptions());
        parser.setKind(ASTParser.K_COMPILATION_UNIT);

        // resolve with root path, so we got the unit name
        String unitName =
            sourceFile.getAbsolutePath().replace(sourceRoot.getParentFile().getAbsolutePath(), "");
        parser.setUnitName(unitName);
        parser.setSource(readFileToString(sourceFile).toCharArray());
        CompilationUnit compilationUnit = (CompilationUnit) parser.createAST(null);
        assert compilationUnit.getAST().hasBindingsRecovery();

        return compilationUnit;
    }

    // returns {(source, source root)}
    HashMap<File, File> listJavaFileMap(String[] sources, boolean removeTests, boolean recursive) {
        HashMap<File, File> files = new HashMap<>();
        for (String sourceRootPath : sources) {
            File sourceRoot = new File(sourceRootPath);
            for (File sourceFile :
                FileUtils.listFiles(sourceRoot, new String[] {"java"}, recursive)) {
                if (removeTests && sourceFile.getName().contains("Test")) {
                    continue;
                }
                files.put(sourceFile, sourceRoot);
            }
        }
        return files;
    }

    List<File> listJavaFiles(String[] sources, boolean removeTests, boolean recursive) {
        List<File> files = new ArrayList<>();
        for (String sourceRootPath : sources) {
            File sourceRoot = new File(sourceRootPath);
            if (removeTests) {
                files.addAll(FileUtils.listFiles(sourceRoot, new String[] {"java"}, recursive)
                                 .stream()
                                 .filter(f -> !removeTests || !f.getName().contains("Test"))
                                 .collect(Collectors.toSet()));
            }
            files.addAll(FileUtils.listFiles(sourceRoot, new String[] {"java"}, recursive));
        }
        return files;
    }

    Set<UncheckedExceptionRelatedMethod> retrieveUncheckedExceptionRelatedMethodsFrom(
        CompilationUnit compilationUnit, Set<String> uncheckedExceptions) {
        UncheckedExceptionRelatedMethodVisitor visitor =
            new UncheckedExceptionRelatedMethodVisitor(uncheckedExceptions);
        compilationUnit.accept(visitor);
        return visitor.getRelatedMethods();
    }

    public Set<UncheckedExceptionRelatedMethod> scanAll(
        String[] sources, Set<String> uncheckedExceptions) {
        if (uncheckedExceptions == null || uncheckedExceptions.isEmpty()) {
            logger.warn("reject scanning due to empty target unchecked exceptions");
            return null;
        }

        HashSet<UncheckedExceptionRelatedMethod> relatedMethods = new HashSet<>();
        for (Map.Entry<File, File> entry : listJavaFileMap(sources, true, true).entrySet()) {
            try {
                logger.info("compiling and scanning source file {}", entry.getKey());
                CompilationUnit compilationUnit =
                    compileSource(sources, entry.getKey(), entry.getValue());
                relatedMethods.addAll(retrieveUncheckedExceptionRelatedMethodsFrom(
                    compilationUnit, uncheckedExceptions));
            } catch (IOException e) {
                logger.info("failed to compile source file {}", entry.getKey());
                e.printStackTrace();
            } catch (Exception e) {
                logger.info("failed to scan ast of source file {}", entry.getKey());
                e.printStackTrace();
            }
        }

        return relatedMethods;
    }

    public Set<UncheckedExceptionRelatedMethod> scanAllInParallel(
        String[] sources, Set<String> uncheckedExceptions) {
        if (uncheckedExceptions == null || uncheckedExceptions.isEmpty()) {
            logger.warn("reject scanning due to empty target unchecked exceptions");
            return null;
        }

        final ConcurrentHashMap<String, Set<UncheckedExceptionRelatedMethod>> relatedMethodMap =
            new ConcurrentHashMap<>();
        final AtomicInteger count = new AtomicInteger(0);

        // parallelStream, start as many threads as core num
        HashMap<File, File> sourceFiles = listJavaFileMap(sources, true, true);

        sourceFiles.entrySet().parallelStream().forEach(entry -> {
            File sourceFile = entry.getKey();
            try {
                logger.info("compiling and scanning source file {}", entry.getKey());
                CompilationUnit compilationUnit =
                    compileSource(sources, entry.getKey(), entry.getValue());
                Set<UncheckedExceptionRelatedMethod> relatedMethods =
                    retrieveUncheckedExceptionRelatedMethodsFrom(
                        compilationUnit, uncheckedExceptions);

                // store in method map
                relatedMethodMap.put(sourceFile.getAbsolutePath(), relatedMethods);
                count.addAndGet(relatedMethods.size());
            } catch (IOException e) {
                logger.info("failed to compile source file {}", entry.getKey());
                e.printStackTrace();
            } catch (Exception e) {
                logger.info("failed to scan ast of source file {}", entry.getKey());
                e.printStackTrace();
            }
        });

        Set<UncheckedExceptionRelatedMethod> relatedMethods = new HashSet<>(count.get());
        for (Map.Entry<String, Set<UncheckedExceptionRelatedMethod>> entry :
            relatedMethodMap.entrySet()) {
            relatedMethods.addAll(entry.getValue());
        }

        return relatedMethods;
    }

    public Set<UncheckedExceptionRelatedMethod> scanAllInParallelUsingExecutors(
        String[] sources, Set<String> uncheckedExceptions) {
        if (uncheckedExceptions == null || uncheckedExceptions.isEmpty()) {
            logger.warn("reject scanning due to empty target unchecked exceptions");
            return null;
        }

        final ConcurrentHashMap<String, Set<UncheckedExceptionRelatedMethod>> relatedMethodMap =
            new ConcurrentHashMap<>();
        final AtomicInteger count = new AtomicInteger(0);

        // parallelStream, start as many threads as core num
        HashMap<File, File> sourceFiles = listJavaFileMap(sources, true, true);

        ExecutorService executorService = Executors.newFixedThreadPool(4);

        for (Map.Entry<File, File> entry : sourceFiles.entrySet()) {
            executorService.submit(() -> {
                try {
                    logger.info("compiling and scanning source file {}", entry.getKey());
                    File sourceFile = entry.getKey();
                    CompilationUnit compilationUnit =
                        compileSource(sources, entry.getKey(), entry.getValue());
                    Set<UncheckedExceptionRelatedMethod> relatedMethods =
                        retrieveUncheckedExceptionRelatedMethodsFrom(
                            compilationUnit, uncheckedExceptions);

                    // store in method map
                    relatedMethodMap.put(sourceFile.getAbsolutePath(), relatedMethods);
                    count.addAndGet(relatedMethods.size());
                } catch (IOException e) {
                    logger.info("failed to compile source file {}", entry.getKey());
                    e.printStackTrace();
                } catch (Exception e) {
                    logger.info("failed to scan ast of source file {}", entry.getKey());
                    e.printStackTrace();
                }
            });
        }
        executorService.shutdown();

        try {
            executorService.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        Set<UncheckedExceptionRelatedMethod> relatedMethods = new HashSet<>(count.get());
        for (Map.Entry<String, Set<UncheckedExceptionRelatedMethod>> entry :
            relatedMethodMap.entrySet()) {
            relatedMethods.addAll(entry.getValue());
        }

        return relatedMethods;
    }

    public String[] getClasspath() {
        return Arrays.copyOf(classpath, classpath.length);
    }

    public static final class UncheckedExceptionRelatedMethod {
        private String packageName;
        private String className;
        private String methodName;
        private String returnTypeName;
        private String[] typeParameters;
        private String[] parameterTypeName;
        private String[] checkedExceptionNames;
        private String[] uncheckedExceptionNames;
        private JavadocType[] exceptionsFromJavaDoc;

        public UncheckedExceptionRelatedMethod() {}

        @Override
        public int hashCode() {
            return Objects.hash(packageName, className, methodName, returnTypeName,
                parameterTypeName, checkedExceptionNames, uncheckedExceptionNames, typeParameters);
        }

        @Override
        public boolean equals(Object obj) {
            if (!(obj instanceof UncheckedExceptionRelatedMethod)) {
                return false;
            }

            UncheckedExceptionRelatedMethod other = (UncheckedExceptionRelatedMethod) obj;
            return Objects.equals(packageName, other.packageName)
                && Objects.equals(className, other.className)
                && Objects.equals(methodName, other.methodName)
                && Objects.equals(returnTypeName, other.returnTypeName)
                && Arrays.equals(parameterTypeName, other.parameterTypeName)
                && Arrays.equals(checkedExceptionNames, other.checkedExceptionNames)
                && Arrays.equals(uncheckedExceptionNames, other.uncheckedExceptionNames)
                && Arrays.equals(typeParameters, other.typeParameters);
        }

        public String[] getTypeParameters() {
            return typeParameters;
        }

        public void setTypeParameters(String[] typeParameters) {
            this.typeParameters = typeParameters;
        }

        public String getPackageName() {
            return packageName;
        }

        public void setPackageName(String packageName) {
            this.packageName = packageName;
        }

        public String getClassName() {
            return className;
        }

        public void setClassName(String className) {
            this.className = className;
        }

        public String getMethodName() {
            return methodName;
        }

        public void setMethodName(String methodName) {
            this.methodName = methodName;
        }

        public String getReturnTypeName() {
            return returnTypeName;
        }

        public void setReturnTypeName(String returnTypeName) {
            this.returnTypeName = returnTypeName;
        }

        public String[] getParameterTypeName() {
            return parameterTypeName;
        }

        public void setParameterTypeName(String[] parameterTypeName) {
            this.parameterTypeName = parameterTypeName;
        }

        public String[] getCheckedExceptionNames() {
            return checkedExceptionNames;
        }

        public void setCheckedExceptionNames(String[] checkedExceptionNames) {
            this.checkedExceptionNames = checkedExceptionNames;
        }

        public String[] getUncheckedExceptionNames() {
            return uncheckedExceptionNames;
        }

        public void setUncheckedExceptionNames(String[] uncheckedExceptionNames) {
            this.uncheckedExceptionNames = uncheckedExceptionNames;
        }

        public JavadocType[] getExceptionsFromJavaDoc() {
            return exceptionsFromJavaDoc;
        }

        public void setExceptionsFromJavaDoc(JavadocType[] exceptionsFromJavaDoc) {
            this.exceptionsFromJavaDoc = exceptionsFromJavaDoc;
        }

        public static class JavadocType {
            private String tag;
            private String type;

            public JavadocType(String tag, String type) {
                this.tag = tag;
                this.type = type;
            }

            public String getTag() {
                return tag;
            }

            public void setTag(String tag) {
                this.tag = tag;
            }

            public String getType() {
                return type;
            }

            public void setType(String type) {
                this.type = type;
            }
        }
    }

    public static class UncheckedExceptionRelatedMethodVisitor extends ASTVisitor {
        private static final Set<String> NEEDED_TAGS = new HashSet<>();

        static {
            NEEDED_TAGS.addAll(Arrays.asList(
                TagElement.TAG_EXCEPTION, TagElement.TAG_LINK, TagElement.TAG_THROWS));
        }

        private final Logger logger =
            LoggerFactory.getLogger(UncheckedExceptionRelatedMethodVisitor.class);
        private final Set<String> targetUncheckedExceptions;
        private PackageDeclaration currentPackageDeclaration;
        private Stack<TypeDeclaration> typeDeclarationStack = new Stack<>();
        private Stack<MethodDeclaration> methodDeclarationStack = new Stack<>();
        private ArrayList<String> thrownUncheckedExceptions = new ArrayList<>();
        private Set<UncheckedExceptionRelatedMethod> relatedMethods = new HashSet<>();

        public UncheckedExceptionRelatedMethodVisitor(Set<String> targetUncheckedExceptions) {
            this.targetUncheckedExceptions = targetUncheckedExceptions;
        }

        private static String getSimpleName(String qualifiedName) {
            String[] splits = qualifiedName.split("\\.");
            return splits[splits.length - 1];
        }

        @SuppressWarnings("unchecked")
        private Collection<UncheckedExceptionRelatedMethod.JavadocType> getRelatedJavadocTypes(
            TagElement tag) {
            if (NEEDED_TAGS.contains(tag.getTagName())) {
                return targetUncheckedExceptions.stream()
                    .filter(e -> tag.toString().contains(getSimpleName(e)))
                    .map(e -> new UncheckedExceptionRelatedMethod.JavadocType(tag.getTagName(), e))
                    .collect(Collectors.toList());
            }

            List<UncheckedExceptionRelatedMethod.JavadocType> types = new ArrayList<>();
            List<IDocElement> fragments = tag.fragments();
            for (IDocElement e : fragments) {
                if (!(e instanceof TagElement)) {
                    continue;
                }

                types.addAll(getRelatedJavadocTypes((TagElement) e));
            }

            return types;
        }

        private UncheckedExceptionRelatedMethod.JavadocType[] getRelatedJavadocTypes(Javadoc doc) {
            if (doc == null) {
                return null;
            }

            List<UncheckedExceptionRelatedMethod.JavadocType> types = new ArrayList<>();
            for (Object tagObj : doc.tags()) {
                final TagElement tag = (TagElement) tagObj;
                types.addAll(getRelatedJavadocTypes(tag));
            }

            return types.toArray(new UncheckedExceptionRelatedMethod.JavadocType[0]);
        }

        private Set<String> getTypeParameters() {
            Set<String> typeParameters = new HashSet<>();
            for (TypeDeclaration typeDeclaration : typeDeclarationStack) {
                List<TypeParameter> params = typeDeclaration.typeParameters();
                typeParameters.addAll(params.stream()
                                          .map(p -> p.getName().getFullyQualifiedName())
                                          .collect(Collectors.toList()));
            }
            for (MethodDeclaration methodDeclaration : methodDeclarationStack) {
                List<TypeParameter> params = methodDeclaration.typeParameters();
                typeParameters.addAll(params.stream()
                                          .map(p -> p.getName().getFullyQualifiedName())
                                          .collect(Collectors.toList()));
            }
            return typeParameters;
        }

        @SuppressWarnings("unchecked")
        private UncheckedExceptionRelatedMethod newUncheckedExceptionRelatedMethod(
            MethodDeclaration methodDeclaration, List<String> thrownUncheckedExceptions,
            UncheckedExceptionRelatedMethod.JavadocType[] javadocTypes) {
            UncheckedExceptionRelatedMethod relatedMethod = new UncheckedExceptionRelatedMethod();

            // set fields
            relatedMethod.setPackageName(
                currentPackageDeclaration.getName().getFullyQualifiedName());
            relatedMethod.setClassName(
                typeDeclarationStack.peek().getName().getFullyQualifiedName());
            relatedMethod.setMethodName(methodDeclaration.getName().getFullyQualifiedName());
            relatedMethod.setReturnTypeName(
                methodDeclaration.getReturnType2().resolveBinding().getName());

            List<SingleVariableDeclaration> parameters = methodDeclaration.parameters();
            relatedMethod.setParameterTypeName(
                parameters.stream()
                    .map(param -> param.getType().resolveBinding().getQualifiedName())
                    .toArray(String[] ::new));

            List<Type> thrownExceptionTypes = methodDeclaration.thrownExceptionTypes();
            relatedMethod.setCheckedExceptionNames(
                thrownExceptionTypes.stream()
                    .map(t -> t.resolveBinding().getQualifiedName())
                    .toArray(String[] ::new));

            if (!thrownUncheckedExceptions.isEmpty()) {
                relatedMethod.setUncheckedExceptionNames(
                    thrownUncheckedExceptions.toArray(new String[0]));
            }
            relatedMethod.setExceptionsFromJavaDoc(javadocTypes);

            // set type params
            relatedMethod.setTypeParameters(getTypeParameters().toArray(new String[0]));

            return relatedMethod;
        }

        // get current package declaration
        @Override
        public boolean visit(PackageDeclaration node) {
            currentPackageDeclaration = node;
            return true;
        }

        // get current type declaration
        @Override
        public boolean visit(TypeDeclaration node) {
            typeDeclarationStack.push(node);
            return true;
        }

        @Override
        public void endVisit(TypeDeclaration node) {
            typeDeclarationStack.pop();
        }

        // get current method declaration
        @Override
        public boolean visit(MethodDeclaration node) {
            methodDeclarationStack.push(node);
            return true;
        }

        @Override
        public void endVisit(MethodDeclaration node) {
            UncheckedExceptionRelatedMethod.JavadocType[] javadocTypes =
                getRelatedJavadocTypes(methodDeclarationStack.peek().getJavadoc());
            // construct an UncheckedExceptionRelatedMethod if
            //  1. throw unchecked exceptions in statements, or
            //  2. find related unchecked exceptions in Javadoc
            if (!thrownUncheckedExceptions.isEmpty()
                || (javadocTypes != null && javadocTypes.length > 0)) {
                try {
                    relatedMethods.add(newUncheckedExceptionRelatedMethod(
                        methodDeclarationStack.peek(), thrownUncheckedExceptions, javadocTypes));
                } catch (Exception e) {
                    logger.error("failed to create new related method for method {}, due to",
                        methodDeclarationStack.peek().getName().getFullyQualifiedName(),
                        e.getClass().getSimpleName());
                }
            }

            methodDeclarationStack.pop();
            thrownUncheckedExceptions.clear();
        }

        // visit throw statement to see if this is an unchecked exception
        @Override
        public boolean visit(ThrowStatement node) {
            try {
                String exceptionName;
                if (node.getExpression() instanceof MethodInvocation) {
                    exceptionName = ((MethodInvocation) node.getExpression())
                                        .resolveMethodBinding()
                                        .getReturnType()
                                        .getQualifiedName();
                } else {
                    exceptionName = node.getExpression().resolveTypeBinding().getQualifiedName();
                }

                if (targetUncheckedExceptions.contains(exceptionName)) {
                    thrownUncheckedExceptions.add(exceptionName);
                }
            } catch (Exception e) {
                logger.error("failed to get exception type for {}, due to ", node.toString(),
                    e.getClass().getSimpleName());
            }
            return true;
        }

        public Set<UncheckedExceptionRelatedMethod> getRelatedMethods() {
            return relatedMethods;
        }
    }
}
